from typing import Any, Dict, List

from langchain.callbacks.base import BaseCallbackHandler
from langchain_core.outputs import LLMResult
from langchain_ollama import ChatOllama

from configs import app_config
from logger import logger


def get_llm():
    class LoggingHandler(BaseCallbackHandler):
        def on_llm_start(
            self, serialized: Dict[str, Any], prompts: List[str], **kwargs
        ):
            logger.info(f"LLM started, prompts: {prompts[0]}")

        def contains_unicode(self, s):
            return "\\u4" in s

        def on_llm_end(self, response: LLMResult, **kwargs) -> None:
            llm_out = response.generations[0][0].text.strip()
            # if self.contains_unicode(llm_out):
            #     llm_out = llm_out.encode("utf-8").decode("unicode_escape")
            # response.generations[0][0].text = llm_out
            logger.info(f"Chat model ended, response: {llm_out}")

    return ChatOllama(
        model=app_config.OLLAMA_LLM_MODEL,
        base_url=app_config.OLLAMA_LLM_URL,
        callbacks=[LoggingHandler()],
    )


ollama_llm = get_llm()
